import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


# ---------------------- SuperGlue 模型定义 ----------------------
class SuperGlue(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.weights = config['weights']
        self.sinkhorn_iterations = config['sinkhorn_iterations']
        self.match_threshold = config['match_threshold']
        self.load_state_dict(torch.load(f'superglue_{config["weights"]}.pth', map_location='cpu'))
        self.eval()

    def forward(self, data):
        return self._match(data['keypoints0'], data['keypoints1'], data['descriptors0'], data['descriptors1'])

    def _match(self, kpts0, kpts1, desc0, desc1):
        # 简化版匹配逻辑（实际需复杂计算）
        scores = torch.einsum('bcn,bcm->bnm', desc0, desc1)
        return {'matches0': torch.argmax(scores, dim=2)}
